import yaml
from pathlib import Path
from collections import defaultdict


class ElementLoader:
    def __init__(self, device_type):
        self.device_type = device_type
        self.elements = self._load_elements(self.device_type)

    def _load_elements(self, device_type):
        config_subdir = "elements_" + device_type
        elements_dir = Path(__file__).parent.parent / "base" / config_subdir

        if not elements_dir.exists():
            raise FileNotFoundError(f"元素目录不存在: {elements_dir}")

        combined = defaultdict(dict)

        # 递归遍历所有YAML文件
        for yaml_file in elements_dir.rglob("*.yaml"):
            try:
                with open(yaml_file, encoding='utf-8') as f:
                    data = yaml.safe_load(f)

                    # 新增数据有效性检查
                    if not isinstance(data, dict):
                        print(f"警告：文件 {yaml_file} 内容不是字典格式，已跳过")
                        continue

                    for page, elements in data.items():
                        if not isinstance(elements, dict):
                            print(f"警告：文件 {yaml_file} 中页面 {page} 的内容不是字典格式，已跳过")
                            continue
                        combined[page].update(elements)

            except yaml.YAMLError as e:
                print(f"YAML解析错误 ({yaml_file}): {str(e)}")
            except Exception as e:
                print(f"加载文件 {yaml_file} 时发生意外错误: {str(e)}")

        return dict(combined)

    def get_locator(self, page_name, element_name):
        return self.elements.get(page_name, {}).get(element_name, {})